/*
 * Copyright 2024-2025 the original author or authors.
 */

package io.modelcontextprotocol.server;

import io.modelcontextprotocol.spec.McpError;
import io.modelcontextprotocol.spec.McpSchema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;

import java.util.Map;

class DefaultMcpStatelessServerHandler implements McpStatelessServerHandler {

    private static final Logger logger = LoggerFactory.getLogger(DefaultMcpStatelessServerHandler.class);

    Map<String, McpStatelessRequestHandler<?>> requestHandlers;

    Map<String, McpStatelessNotificationHandler> notificationHandlers;

    public DefaultMcpStatelessServerHandler(Map<String, McpStatelessRequestHandler<?>> requestHandlers,
                                            Map<String, McpStatelessNotificationHandler> notificationHandlers) {
        this.requestHandlers = requestHandlers;
        this.notificationHandlers = notificationHandlers;
    }

    @Override
    public Mono<McpSchema.JSONRPCResponse> handleRequest(McpTransportContext transportContext,
                                                         McpSchema.JSONRPCRequest request) {
        McpStatelessRequestHandler<?> requestHandler = this.requestHandlers.get(request.getMethod());
        if (requestHandler == null) {
            return Mono.error(new McpError("Missing handler for request type: " + request.getMethod()));
        }
        return requestHandler.handle(transportContext, request.getParams())
                .map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), result, null))
                .onErrorResume(t -> {
                    McpSchema.JSONRPCResponse.JSONRPCError error = null;
                    if (t instanceof McpError) {
                        McpError mcpError = (McpError) t;
                        if (mcpError.getJsonRpcError() != null) {
                            error = mcpError.getJsonRpcError();
                        }
                    }

                    if (error == null) {
                        error = new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
                                t.getMessage(), null);
                    }

                    return Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.getId(), null, error));
                });
    }

    @Override
    public Mono<Void> handleNotification(McpTransportContext transportContext,
                                         McpSchema.JSONRPCNotification notification) {
        McpStatelessNotificationHandler notificationHandler = this.notificationHandlers.get(notification.getMethod());
        if (notificationHandler == null) {
            logger.warn("Missing handler for notification type: {}", notification.getMethod());
            return Mono.empty();
        }
        return notificationHandler.handle(transportContext, notification.getParams());
    }

}
